package io.gitee.kenewstar.dbmq.data;

import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.date.DateTime;
import cn.hutool.core.date.DateUtil;
import cn.hutool.core.util.RandomUtil;
import io.gitee.kenewstar.dbmq.ConsumerRegisterInfo;
import io.gitee.kenewstar.dbmq.ConsumerTask;
import io.gitee.kenewstar.dbmq.Message;
import io.gitee.kenewstar.dbmq.mongo.MongoPrimaryKey;
import io.gitee.kenewstar.dbmq.util.Const;
import io.gitee.kenewstar.dbmq.mongo.MongoKey;
import io.gitee.kenewstar.dbmq.util.Util;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.data.mongodb.core.BulkOperations;
import org.springframework.data.mongodb.core.FindAndModifyOptions;
import org.springframework.data.mongodb.core.MongoTemplate;
import org.springframework.data.mongodb.core.query.Criteria;
import org.springframework.data.mongodb.core.query.Query;
import org.springframework.data.mongodb.core.query.Update;

import java.util.ArrayList;
import java.util.List;
import java.util.Objects;
import java.util.Optional;

/**
 * @author kenewstar
 */
public class DbMongoHandler implements DbDataHandler, DbDataHandlerAdapter {

    private static final Logger LOGGER = LoggerFactory.getLogger(DbMongoHandler.class);

    private MongoTemplate mongoTemplate;

    public DbMongoHandler() {}

    public DbMongoHandler(MongoTemplate mongoTemplate) {
        this.mongoTemplate = mongoTemplate;
    }

    @Override
    public DbDataHandler adapter() {
        return this;
    }

    @Override
    public String getType() {
        return DbDataHandlerWrapper.MONGO;
    }

    private Long getPrimaryKey(String table) {
        Query query = new Query(Criteria.where(MongoKey.ID).is(table));
        Update update = new Update().inc(MongoKey.VALUE, 1);
        final MongoPrimaryKey primaryKey = this.mongoTemplate.findAndModify(query, update,
                FindAndModifyOptions.options().upsert(true).returnNew(true), MongoPrimaryKey.class);
        return Optional.ofNullable(primaryKey).orElse(new MongoPrimaryKey()).getValue();
    }

    private <T> T atomicUpdate(Query query, Update update, Class<T> clz, String table) {
        return this.mongoTemplate.findAndModify(query, update,
                FindAndModifyOptions.options().upsert(false).returnNew(true),
                clz, table);
    }

    @Override
    public void initConsumerRegister() {
        LOGGER.info("-------------------no operation---------------------------");
    }

    @Override
    public void initConsumerQueue(List<ConsumerTask> taskList) {
        LOGGER.info("-------------------no operation---------------------------");
    }

    @Override
    public List<ConsumerRegisterInfo> pullRegisterInfo(List<String> topics) {
        if (CollUtil.isEmpty(topics)) {
            return CollUtil.newArrayList();
        }
        Query query = new Query(Criteria.where(ConsumerRegisterInfo.TOPIC).in(topics));
        return this.mongoTemplate.find(query, ConsumerRegisterInfo.class, Const.Table.MQ_CONSUMER_REGISTER);
    }

    @Override
    public ConsumerRegisterInfo createConsumer(ConsumerTask task) {
        ConsumerRegisterInfo info = new ConsumerRegisterInfo();
        info.setId(this.getPrimaryKey(Const.Table.MQ_CONSUMER_REGISTER));
        info.setTopic(task.getTopic());
        info.setUuid(task.getUuid());
        info.setHeartbeatTime(DateUtil.date());
        info.setIp(Util.getIpAddress());
        info.setConsumeCount(0L);
        info.setStatus(Const.Status.UP);
        info.setVersion(1L);
        this.mongoTemplate.insert(info, Const.Table.MQ_CONSUMER_REGISTER);
        task.setId(info.getId());
        return info;
    }

    @Override
    public void updateConsumer(ConsumerTask task, List<ConsumerRegisterInfo> registerInfo) {

        Update update = new Update().set(ConsumerRegisterInfo.HEARTBEAT_TIME, DateUtil.date())
                .set(ConsumerRegisterInfo.IP, Util.getIpAddress())
                .set(ConsumerRegisterInfo.UUID, task.getUuid())
                .set(ConsumerRegisterInfo.STATUS, Const.Status.UP)
                .inc(ConsumerRegisterInfo.VERSION, 1);

        ConsumerRegisterInfo res = null;
        for (ConsumerRegisterInfo info : registerInfo) {
            Query query = new Query(Criteria.where(MongoKey.ID).is(info.getId())
                    .and(ConsumerRegisterInfo.UUID).is(info.getUuid())
                    .and(ConsumerRegisterInfo.VERSION).is(info.getVersion()));
            res = this.atomicUpdate(query, update, ConsumerRegisterInfo.class, Const.Table.MQ_CONSUMER_REGISTER);
            // 执行成功
            if (Objects.nonNull(res)) {
                task.setOldId(task.getId());
                task.setId(res.getId());
                break;
            }
        }
        if (Objects.isNull(res)) {
            this.createConsumer(task);
        }

    }

    @Override
    public void registerDown(Long id) {
        if (Objects.isNull(id)) {
            return;
        }
        Update update = new Update().set(ConsumerRegisterInfo.STATUS, Const.Status.DOWN);
        Query query = new Query(Criteria.where(MongoKey.ID).is(id));
        this.atomicUpdate(query, update, ConsumerRegisterInfo.class, Const.Table.MQ_CONSUMER_REGISTER);
    }

    @Override
    public void registerDown(List<ConsumerRegisterInfo> list) {
        if (CollUtil.isEmpty(list)) {
            return;
        }
        Update update = new Update().set(ConsumerRegisterInfo.STATUS, Const.Status.DOWN).inc(ConsumerRegisterInfo.VERSION, 1);
        for (ConsumerRegisterInfo info : list) {
            Query query = new Query(Criteria.where(MongoKey.ID).is(info.getId())
                    .and(ConsumerRegisterInfo.STATUS).is(Const.Status.UP)
                    .and(ConsumerRegisterInfo.VERSION).is(info.getVersion())
                    .and(ConsumerRegisterInfo.HEARTBEAT_TIME).is(info.getHeartbeatTime())
            );
            this.atomicUpdate(query, update, ConsumerRegisterInfo.class, Const.Table.MQ_CONSUMER_REGISTER);
        }
    }

    @Override
    public void sendHeartbeat(List<Long> ids) {
        if (CollUtil.isEmpty(ids)) {
            return;
        }

        LOGGER.info("=====================Send Heartbeat===========================");
        Update update = new Update().set(ConsumerRegisterInfo.STATUS, Const.Status.UP).set(ConsumerRegisterInfo.HEARTBEAT_TIME, DateUtil.date());
        Query query = new Query(Criteria.where(MongoKey.ID).in(ids));
        this.atomicUpdate(query, update, ConsumerRegisterInfo.class, Const.Table.MQ_CONSUMER_REGISTER);
    }

    @Override
    public List<Message> getMessage(ConsumerTask task) {
        if (Objects.isNull(task.getTag())) {
            return CollUtil.newArrayList();
        }
        Query query = new Query(Criteria
                .where(Message.CONSUME_FLAG).is(0)
                .and(Message.CONSUMER_TAG).is(task.getTag())
        ).limit(task.getLimit());
        return this.mongoTemplate.find(query, Message.class, task.getTopic());
    }

    @Override
    public void deleteMessage(String topic, List<Long> ids) {
        if (CollUtil.isEmpty(ids)) {
            return;
        }
        Query query = new Query(Criteria.where(MongoKey.ID).in(ids));
        this.mongoTemplate.findAllAndRemove(query, topic);
    }

    @Override
    public int sendMessage(String topic, List<String> messageBodyList, List<Integer> tags) {
        final DateTime current = DateUtil.date();
        List<Message> list = new ArrayList<>(messageBodyList.size());
        for (String messageBody : messageBodyList) {
            Message message = new Message();
            message.setId(this.getPrimaryKey(topic));
            message.setMessageBody(messageBody);
            message.setConsumeFlag(0);
            if (CollUtil.isNotEmpty(tags)) {
                message.setConsumerTag(tags.get(RandomUtil.randomInt(tags.size())));
            } else {
                message.setConsumerTag(1);
            }
            message.setCreationDate(current);
            list.add(message);
        }
        return this.mongoTemplate.insert(list, topic).size();
    }

    @Override
    public List<ConsumerRegisterInfo> getRegisterInfo(List<String> topics) {
        Query query = new Query(Criteria.where(ConsumerRegisterInfo.TOPIC).in(topics));
        return this.mongoTemplate.find(query, ConsumerRegisterInfo.class, Const.Table.MQ_CONSUMER_REGISTER);
    }

    @Override
    public boolean tryGetConsumerRegister(ConsumerRegisterInfo registerInfo) {
        if (!Objects.equals(registerInfo.getStatus(), Const.Status.DOWN)) {
            return Boolean.FALSE;
        }
        Update update = new Update().set(ConsumerRegisterInfo.STATUS, Const.Status.CLEAR);
        Query query = new Query(Criteria.where(MongoKey.ID).is(registerInfo.getId()).and(ConsumerRegisterInfo.STATUS).is(Const.Status.DOWN));
        final ConsumerRegisterInfo updateObj = this.atomicUpdate(query, update, ConsumerRegisterInfo.class, Const.Table.MQ_CONSUMER_REGISTER);
        return Objects.nonNull(updateObj);
    }

    @Override
    public void releaseConsumerRegister(ConsumerRegisterInfo registerInfo) {
        Update update = new Update().set(ConsumerRegisterInfo.STATUS, Const.Status.DOWN);
        Query query = new Query(Criteria.where(MongoKey.ID).is(registerInfo.getId()).and(ConsumerRegisterInfo.STATUS).is(Const.Status.CLEAR));
        this.atomicUpdate(query, update, ConsumerRegisterInfo.class, Const.Table.MQ_CONSUMER_REGISTER);
    }

    @Override
    public void deleteConsumerRegister(ConsumerRegisterInfo registerInfo) {
        Query query = new Query(Criteria.where(MongoKey.ID).is(registerInfo.getId()).and(ConsumerRegisterInfo.STATUS).is(Const.Status.CLEAR));
        this.mongoTemplate.remove(query, Const.Table.MQ_CONSUMER_REGISTER);
    }

    @Override
    public List<Message> getReAssignMessage(String topic, Integer tag) {
        Query query = new Query(Criteria.where(Message.CONSUMER_TAG).is(tag)
                .and(Message.CONSUME_FLAG).is(0)
        ).limit(COUNT);
        return this.mongoTemplate.find(query, Message.class, topic);
    }

    @Override
    public void batchUpdateMessage(List<Message> messageList, List<Integer> tags, String topic) {
        final BulkOperations bulkOperations = this.mongoTemplate.bulkOps(BulkOperations.BulkMode.UNORDERED, topic);
        for (Message message : messageList) {
            Integer newTag = Util.getNewTag(tags);
            Update update = new Update().set(Message.CONSUMER_TAG, newTag);
            bulkOperations.updateOne(new Query(Criteria.where(MongoKey.ID).is(message.getId())), update);
        }
        bulkOperations.execute();
    }
}
